package cn.ddiancan.auth.service.provider;

import java.security.Principal;
import java.util.Base64;
import java.util.Objects;
import java.util.Set;

import cn.ddiancan.auth.service.token.XddAuth2GrantAuthenticationToken;
import cn.ddiancan.auth.service.token.XddAuth2UserCodeAuthenticationToken;
import cn.ddiancan.auth.utils.JwtUtils;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.crypto.keygen.Base64StringKeyGenerator;
import org.springframework.security.crypto.keygen.StringKeyGenerator;
import org.springframework.security.oauth2.core.AuthorizationGrantType;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2ErrorCodes;
import org.springframework.security.oauth2.core.OAuth2RefreshToken;
import org.springframework.security.oauth2.core.OAuth2Token;
import org.springframework.security.oauth2.core.OAuth2UserCode;
import org.springframework.security.oauth2.core.endpoint.OAuth2ParameterNames;
import org.springframework.security.oauth2.jwt.JwsHeader;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.jwt.JwtEncoder;
import org.springframework.security.oauth2.jwt.JwtEncoderParameters;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AccessTokenAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AuthorizationGrantAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContextHolder;
import org.springframework.security.oauth2.server.authorization.settings.AuthorizationServerSettings;
import org.springframework.security.oauth2.server.authorization.token.DefaultOAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;

public class OAuth2AccessTokenAuthenticationTokenHelper {

    private static final String ACTUAL_TOKEN = "actualToken";
    private static final String REQUEST_SOURCE = "requestSource";
    private static final String LOGIN_TYPE = "loginType";

    private static final StringKeyGenerator DEFAULT_TOKEN_GENERATOR =
        new Base64StringKeyGenerator(Base64.getUrlEncoder().withoutPadding(), 96);

    public static OAuth2AccessTokenAuthenticationToken auth2AccessTokenAuthenticationToken(
        XddAuth2GrantAuthenticationToken authentication, UserDetails userDetails,
        AuthorizationServerSettings authorizationServerSettings,
        OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer, OAuth2AuthorizationService oAuth2AuthorizationService,
        OAuth2TokenGenerator<?> tokenGenerator,
        JwtEncoder jwtEncoder) {
        OAuth2ClientAuthenticationToken clientPrincipal =
            OAuth2AccessTokenAuthenticationTokenHelper.getAuthenticatedClientElseThrowInvalidClient(authentication);
        RegisteredClient registeredClient = clientPrincipal.getRegisteredClient();
        // 构造认证信息
        Authentication principal =
            new UsernamePasswordAuthenticationToken(userDetails.getUsername(), userDetails.getPassword(),
                userDetails.getAuthorities());

        OAuth2Authorization authorization =
            OAuth2Authorization.withRegisteredClient(Objects.requireNonNull(clientPrincipal.getRegisteredClient()))
                .principalName(principal.getName()).authorizationGrantType(authentication.getGrantType())
                .attribute(Principal.class.getName(), principal).attribute("scopes", registeredClient.getScopes())
                .build();

        String issuer = authorizationServerSettings != null ? authorizationServerSettings.getIssuer() : null;
        Set<String> authorizedScopes = authorization.getAttribute("scopes");
        // 构造jwt token信息
        JwsHeader.Builder headersBuilder = JwtUtils.headers();
        headersBuilder.header("client-id", registeredClient.getClientId());
        headersBuilder.header("authorization-grant-type", authentication.getGrantType().getValue());
        JwtClaimsSet.Builder claimsBuilder =
            JwtUtils.accessTokenClaims(registeredClient, issuer, authorization.getPrincipalName(), authorizedScopes);
        JwtEncodingContext context =
            JwtEncodingContext.with(headersBuilder, claimsBuilder).registeredClient(registeredClient)
                .principal(authorization.getAttribute(Principal.class.getName())).authorization(authorization)
                .authorizedScopes(authorizedScopes).tokenType(OAuth2TokenType.ACCESS_TOKEN)
                .authorizationGrantType(AuthorizationGrantType.PASSWORD).authorizationGrant(authentication).build();

        jwtCustomizer.customize(context);
        JwsHeader headers = context.getJwsHeader().build();
        JwtClaimsSet claims = context.getClaims().build();
        DefaultOAuth2TokenContext.Builder tokenContextBuilder =
            DefaultOAuth2TokenContext.builder().registeredClient(registeredClient)
                .principal(authorization.getAttribute(Principal.class.getName()))
                .authorizationServerContext(AuthorizationServerContextHolder.getContext())
                .tokenType(OAuth2TokenType.REFRESH_TOKEN).authorization(authorization)
                .authorizedScopes(authorizedScopes).authorizationGrant(authentication);
        // 真实token
        OAuth2AccessToken accessToken = getOAuth2AccessToken(headers, claims, authorizedScopes,jwtEncoder);
        // 虚拟映射token
        String virtualToken = DEFAULT_TOKEN_GENERATOR.generateKey();
        OAuth2AccessToken virtualAccessToken =
            new OAuth2AccessToken(accessToken.getTokenType(), virtualToken, accessToken.getIssuedAt(),
                accessToken.getExpiresAt(), authorizedScopes);
        OAuth2Token generateRefreshToken =
            tokenGenerator.generate(tokenContextBuilder.tokenType(OAuth2TokenType.REFRESH_TOKEN).build());
        OAuth2Authorization oAuth2Authorization =
            OAuth2Authorization.from(authorization).accessToken(virtualAccessToken)
                .refreshToken((OAuth2RefreshToken) generateRefreshToken).attribute(OAuth2ParameterNames.STATE, "1")
                .authorizedScopes(authorizedScopes).attribute(ACTUAL_TOKEN, accessToken.getTokenValue())
                .attribute(REQUEST_SOURCE,authentication.getRequestSource().name())
                .attribute(LOGIN_TYPE,authentication.getRequestType()).build();
        if(authentication instanceof XddAuth2UserCodeAuthenticationToken userCodeAuthenticationToken){
            OAuth2UserCode userCode = new OAuth2UserCode(userCodeAuthenticationToken.getUserCode(),
                accessToken.getIssuedAt(),accessToken.getExpiresAt());
            oAuth2Authorization = OAuth2Authorization.from(oAuth2Authorization).token(userCode).build();
        }
        oAuth2AuthorizationService.save(oAuth2Authorization);
        return new OAuth2AccessTokenAuthenticationToken(registeredClient, clientPrincipal, virtualAccessToken,
            (OAuth2RefreshToken) generateRefreshToken);
    }

    private static OAuth2AccessToken getOAuth2AccessToken(JwsHeader headers, JwtClaimsSet claims,
        Set<String> authorizedScopes,JwtEncoder jwtEncoder) {
        JwtEncoderParameters params = JwtEncoderParameters.from(headers, claims);
        Jwt jwtAccessToken = jwtEncoder.encode(params);
        // 生成token
        return new OAuth2AccessToken(OAuth2AccessToken.TokenType.BEARER, jwtAccessToken.getTokenValue(),
            jwtAccessToken.getIssuedAt(), jwtAccessToken.getExpiresAt(), authorizedScopes);
    }

    static OAuth2ClientAuthenticationToken getAuthenticatedClientElseThrowInvalidClient(Authentication authentication) {
        OAuth2ClientAuthenticationToken clientPrincipal = null;
        if (OAuth2ClientAuthenticationToken.class.isAssignableFrom(authentication.getPrincipal().getClass())) {
            clientPrincipal = (OAuth2ClientAuthenticationToken) authentication.getPrincipal();
        }
        if (clientPrincipal != null && clientPrincipal.isAuthenticated()) {
            return clientPrincipal;
        }
        throw new OAuth2AuthenticationException(OAuth2ErrorCodes.INVALID_CLIENT);
    }

}
